# PDFDecryper-4.5
# Auther:tianwenxiao
# mail:tianwenxiao@wo.cn
# 主要功能:
#        2.0：针对没有设置打开密码的PDF文件，只是设置了打印、编辑等限制的密码，可以使用该脚本进行解密
#        3.5: 新增：针对设置了打开密码的文件，增加了密码字典库爆破功能，只需要选择需要破解的PDF文件，再选择用于破解的字典库，点击执行脚本即可开始破解
#        4.0：新增：针对需要使用密码字典库进行爆破的操作，新增了多线程的方式，通过生成器将字典文件分成块，使用多线程分别遍历，极大的提高了破解效率，但是
#            目前仍然有个问题，那就是目前GUI上的进度条目前无法更新，后续版本即将修复。
#        4.5: 修复：针对4.0版本中遇到的两个问题：①多线程运行时进程阻塞导致无法中途停止的问题②多线程运行时由于进程阻塞导致的进度条不更新的问题，
#                   这两个问题目前通通解决啦，现在可以爽快的进行密码字典爆破，并且可以显示进度，随时打断
#             新增：针对程序运行中出现的异常信息，增加了一个标签显示
from shlex import join
from socket import timeout
import threading
import tkinter as tk
from tkinter import filedialog
#pdf的读取方法
from PyPDF2 import PdfReader
#pdf的写入方法
from PyPDF2 import PdfWriter
#高加密的方法
from Crypto.Cipher import AES
from tkinter import ttk
import concurrent.futures
import queue

from pandas import value_counts


# 定义一个类在线程中共享信息
class SharedState:
    def __init__(self):
        self.state = True
        self.lock = threading.Lock()

    def set_state(self, new_state):
        with self.lock:
            self.state = new_state

    def get_state(self):
        with self.lock:
            return self.state

# 创建一个线程间信息通信对象
thread_staus = SharedState()
# 创建一堆全局变量，用于控制杂七杂八的进程
complete_num = 0
comput_status = True
total_tasks = 0
input_crack_dic_var = None
progress_status = True
# 切分密码库文件的生成器
def chunks_generator(filename, chunk_size):
    with open(filename, 'r', encoding='utf-8') as file:
        chunk = []
        for line in file:
            line = line.strip()
            chunk.append(line)
            if len(chunk) == chunk_size:  # 修改这里
                yield chunk
                chunk = []  # 清空chunk以便收集下一个数据块
        if chunk:  # 处理文件末尾不完整的块
            yield chunk

# 更新进度条的方法
def check_progress(root,progress, progress_label, info_label, progress_queue, info_queue,input_crack_dic_var):
    global complete_num
    global total_tasks
    global comput_status
    global progress_status
    # 实时监控选中的字典路径
    crack_dir = input_crack_dic_var.get()
    try:
        current_status = thread_staus
        if current_status:
            # 如果用户没有选择密码字典，就什么也不做
            if crack_dir == None or crack_dir == "":
                pass
            else:
                # 获取用户选择的字典中总共有多少条密码，这个只会获取一次
                if comput_status:
                    total_tasks = count_lines_in_file(crack_dir)
                    comput_status = False
                progress['max'] = total_tasks
                # print("total_tasks:{}".format(total_tasks))
                # print("共需遍历{}条数据".format(total_tasks))
                # 只有当队列中有数据时，才读取并更新进度条
                if not progress_queue.empty():
                    # print("complete_num==>{}".format(complete_num))
                    if complete_num < total_tasks:
                        completed_work = progress_queue.get(block=True, timeout=5)
                        complete_num = complete_num + completed_work
                        progress['value'] = complete_num
                        progress_label['text'] = "已完成：{:.1f}%".format((progress["value"]/total_tasks)*100)
                    else:
                        info_label['text'] = "遍历了{}条密码，未找到匹配的密码，请换个密码字典试一试吧!".format(total_tasks)
                        progress_status = False
                    root.update_idletasks()
                    progress_queue.task_done()
            if not info_queue.empty():
                info_label['text'] = info_queue.get(block=True, timeout=5)
    except Exception as e:
        info_label['text'] = "出现了错误，请重新运行程序"
        pass
    finally:
        if progress_status:
            root.after(2000,lambda:check_progress(root,progress, progress_label, info_label,progress_queue,info_queue,input_crack_dic_var))

# 统计密码库总共有多少行
def count_lines_in_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return sum(1 for _ in file)

# 选择要处理的文件
def select_input_file(input_file_path_var):
    input_file_path = filedialog.askopenfilename(filetypes=[("PDF files", "*.PDF;*.pdf")])
    if input_file_path:
        input_file_path_var.delete(0, tk.END)
        input_file_path_var.insert(0, input_file_path)
    return input_file_path

# 选择用于破解的密码字典
def select_cracker_dic(input_crack_dic_var):
    crack_dic = filedialog.askopenfilename(filetypes=[("TXT file","*.txt;*.TXT")])
    if crack_dic:
        input_crack_dic_var.delete(0,tk.END)
        input_crack_dic_var.insert(0,crack_dic)
        return crack_dic

# 定义一个函数来读取PDF文件
def get_reader(filename,password=None):
    try:
        old_file = open(filename, 'rb')
    except Exception as err:
        return print('文件打开失败！' + str(err))
    #如果使用python2需要将PdfReader改为PdfFileReader
    pdf_reader = PdfReader(old_file, strict=False)
    # 执行解密操作
    if pdf_reader.is_encrypted:
        if password is None:
            return pdf_reader
        else:
            success = pdf_reader.decrypt(password)
            if not success :
                #return print('密码不正确！--{}'.format(filename))
                return
            else:
                #print("解密成功，密码为{}".format(password))
                return pdf_reader
    elif old_file in locals():
        old_file.close()
        # 返回结果
    return pdf_reader

# 真正处理PDF的方法
def deception_pdf(root,filename,progress_queue,info_queue,count_lines=None,chunk=None) :

    # 如果用户没有选择破解字典库
    if chunk is None or chunk == '' or count_lines is None or count_lines == '' or progress_queue == None:
        #print("==没有选择破解字典库==")
        try:
            pdf_reader = get_reader(filename)
            if pdf_reader is None:
                info_queue.put("无法读取PDF内容，请重试！")
            # 如果使用python2需要将is_encrypted改为isEncrypted
            elif not pdf_reader.is_encrypted:
                info_queue.put("文件没有被加密，无需操作！")
            else:
                # 如果使用的是python2需要将PdfWriter改为PdfFileWriter
                pdf_writer = PdfWriter()
                #如果使用的是python2需要将将append_pages_from_reader改为appendPagesFromReader
                pdf_writer.append_pages_from_reader(pdf_reader)
                #创建解密后的pdf文件和展示文件的路径
                decrypted_filename = "".join(filename.split('.')[:-1]) + '_' + '已解密' + '.pdf'
                info_queue.put("解密文件已生成:{}".format(decrypted_filename))
                # 写入新文件
                pdf_writer.write(open(decrypted_filename, 'wb'))
                current_state = False
                thread_staus.set_state(False)
        except Exception as err:
            info_queue.put("文件设置了打开密码，无法读取PDF，请尝试选择字典库进行爆破")
            current_state = False
            thread_staus.set_state(False)
            root.destroy()
    else :
        current_thread = threading.current_thread()
        thread_name = current_thread.name
        count = 0
        complete_num = 0
        for password in chunk:
                current_state = thread_staus.get_state()
                if current_state:
                    pdf_reader = get_reader(filename, password)
                    count = count + 1
                    complete_num = complete_num + 1
                    if count >= 500:
                        progress_queue.put(count)
                        count = 0
                    # 如果破解成功
                    if(pdf_reader):
                        # 如果使用的是python2需要将PdfWriter改为PdfFileWriter
                        pdf_writer = PdfWriter()
                        #如果使用的是python2需要将将append_pages_from_reader改为appendPagesFromReader
                        pdf_writer.append_pages_from_reader(pdf_reader)
                        #创建解密后的pdf文件和展示文件的路径
                        decrypted_filename = "".join(filename.split('.')[:-1]) + '_' + '已解密' + '.pdf'
                        # 写入新文件
                        pdf_writer.write(open(decrypted_filename, 'wb'))
                        thread_staus.set_state(False)
                        current_state = False
                        info_queue.put("解密完成", "您的PDF文件已破解完成，密码为{}".format(password))
                        break
                    else:
                        if(complete_num <= count_lines):
                            #print(("线程{}正在解析密码==>{}".format(thread_name,password)))
                            pass
                        else:
                            progress_queue.put(count)
                            info_queue.put("抱歉,未找到匹配的密码，请换个字典库再试一下吧")
                            current_state = False
                            thread_staus.set_state(False)
                            break
        progress_queue.put(count)
        count = 0

# 取消破解，退出应用
def closeApp(root):
    thread_staus.set_state(False)
    root.destroy()

# 开启线程，开始破解
def start_deception_pdf_thread(root, filename,progress_queue,info_queue,count_lines,chunks):
        # 使用 ThreadPoolExecutor 管理线程池
        with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
            # 将生成器的每一行作为任务提交到线程池
            futures = [executor.submit(deception_pdf,root,filename, progress_queue,info_queue,count_lines,chunk)for chunk in chunks]
            # 等待所有任务完成（可选）使用这个会导致在遍历完后，程序卡死
            #concurrent.futures.wait(futures)
            # 确保所有进度更新都已完成
            progress_queue.join()

def start_back_threads(root, filename,progress_queue,info_queue,crack_dic=None):
    # 重置终止标志
    thread_staus.set_state(True)
    if crack_dic == None or crack_dic == "":
        deception_pdf(root,filename,progress_queue,info_queue)
    else:
        # 获取字典库总行数，并传到函数里面，用于显示进度条
        count_lines = count_lines_in_file(crack_dic)
        chunks = chunks_generator(crack_dic,1000)
        decryp_pdf_thread = threading.Thread(target=start_deception_pdf_thread,args=(root, filename, progress_queue,info_queue,count_lines,chunks))
        decryp_pdf_thread.start()

# 创建主函数
def main():
    global input_crack_dic_var
    # 创建主窗口
    root = tk.Tk()
    root.title("PDF解密小工具")
    root.geometry("600x300")
    root.resizable(False, False)
    # 创建一个线程安全的队列来传递进度信息
    progress_queue = queue.Queue()
    # 创建一个用于传递状态信息的队列
    info_queue = queue.Queue()
    # 输入文件框
    input_frame = tk.Frame(root)
    input_frame.grid(column=0, row=0, padx=0, pady=10,sticky="w")
    input_label = tk.Label(input_frame, text="输入文件:")
    input_label.grid(column=0, row=0, padx=17, pady=10,sticky="w")
    input_file_path_var = tk.Entry(input_frame, width=50)
    input_file_path_var.grid(column=1, row=0, padx=10, pady=10,sticky="e")
    input_button = tk.Button(input_frame, text="选择文件", command=lambda:select_input_file(input_file_path_var))
    input_button.grid(column=2, row=0, padx=6, pady=10,sticky="e")
    # 选择密码字典
    dic_frame = tk.Frame(root)
    dic_frame.grid(column=0, row=1, padx=0, pady=10,sticky="w")
    dic_frame_label = tk.Label(dic_frame, text="请选择密码字典:")
    dic_frame_label.grid(column=1, row=1, padx=0, pady=10,sticky="w")
    input_crack_dic_var = tk.Entry(dic_frame, width=50)
    input_crack_dic_var.grid(column=2, row=1, padx=10, pady=10,sticky="e")
    dic_frame_button = tk.Button(dic_frame, text="选择文件", command=lambda:select_cracker_dic(input_crack_dic_var))
    dic_frame_button.grid(column=3, row=1, padx=5, pady=10,sticky="e")
    # 创建进度条
    progress_frame= tk.Frame(root)
    progress_frame.grid(column=0, row=3, padx=5, pady=10,sticky="w")
    progress = ttk.Progressbar(progress_frame, style='custom.Horizontal.TProgressbar', orient='horizontal', length=460, mode='determinate')

    # 创建并配置进度条样式
    style = ttk.Style()
    # 定义进度条的样式（这里以设置背景色为例）
    style.configure('custom.Horizontal.TProgressbar', background='blue')
    progress.grid(column=0, row=3, padx=0, pady=10)
    progress_label = tk.Label(progress_frame, text="已完成0%")
    progress_label.grid(column=1, row=3,padx=10, pady=10,sticky="w")
    controller_frame=tk.Frame(root)
    controller_frame.grid(column=0, row=5, padx=0, pady=10)
    info_frame = tk.Frame(root)
    info_frame.grid(column=0, row=4, padx=5, pady=5)
    info_label = tk.Label(info_frame, text="欢迎使用PDFDecryper")
    info_label.grid(column=0, row=4,padx=5, pady=5)

    # 创建按钮用于执行脚本
    execute_button = tk.Button(controller_frame, text="执行脚本", command=lambda:start_back_threads(root,input_file_path_var.get(),progress_queue,info_queue,input_crack_dic_var.get()))
    execute_button.grid(column=0, row=5, padx=10, pady=10)
    exit_button = tk.Button(controller_frame,text="取消并退出",command=lambda:closeApp(root))
    exit_button.grid(column=1, row=5, padx=10, pady=10)
    # 运行Tkinter事件循环
    root.after(1000,lambda:check_progress(root,progress, progress_label, info_label,progress_queue,info_queue,input_crack_dic_var))
    root.mainloop()

if __name__ == "__main__":
    main()
